import pywencai
import pandas as pd

def interactive_stock_query():
    # 用户交互模块
    query = input("请输入选股条件：\n示例：'市值>100亿；市盈率<30'\n请输入：")
    
    try:
        # 核心查询配置[2,3](@ref)
        result = pywencai.get(
            question=query,
            loop=True,        # 自动翻页获取全部数据
            perpage=100,      # 每页最大数据量（问财服务端限制）
            sleep=2,          # 请求间隔防止封IP
            retry=5,          # 失败重试次数
            log=False         # 关闭调试日志
        )
        
        # 数据清洗模块
        if not result.empty:

            # 结果展示与保存
            print(f"共找到{len(result)}条记录，前10条预览：")
            print(result.head(10).to_markdown(index=False))
            
            # 保存完整结果
            save_choice = input("是否保存完整结果？(y/n): ")
            if save_choice.lower() == 'y':
                result.to_excel(f"./stocks_selection/选股结果_{query[:20]}.xlsx", index=False)
                print("文件已保存至程序目录")
        else:
            print("未找到符合条件的股票")
            
    except Exception as e:
        print(f"查询失败，错误类型：{type(e).__name__}")
        print("可能原因：1.网络连接异常 2.查询语法错误 3.服务端限制")

if __name__ == "__main__":
    interactive_stock_query()